#include "refine/smooth_d.h"
#include "refine/window_pass.h"
#include <limits>
#include <iostream>
#include "imglib/imageclass.h"
#include <imglib/bitmap.h>
#include <imglib/bmpmorph.h>
#include <glib/g3/vector.h>

#define GSL_RANGE_CHECK_OFF
#define HAVE_INLINE
#include <gsl/gsl_linalg.h>

using std::clog;
using std::endl;
using std::flush;

namespace refine {

namespace smooth_d {

const params default_params = {2, 2, 1, 100, 100, false};

typedef grey_mask::mask::elem_type grey_mask_value;
typedef unsigned uint;
typedef unsigned long long ulong;
typedef g3::basic_vector<uint>		vectoru;
typedef g3::basic_vector<ulong>		vectorul;
typedef g3::basic_vector<float>		vectorf;

const uint uchar_max = std::numeric_limits<unsigned char>::max();
const uint uchar_avg = uchar_max / 2;

template <typename Color>
inline vectoru v(Color const & c)
	{ return vectoru(c.b, c.g, c.r); }

// reconstructs foreground and background in a point and from them and image color
// updates alpha in that point
template <typename Raster, typename Matte>
class smoother
{
public:
	smoother(Raster const & src_, 
		Matte const & matte_, grey_mask::mask & next_matte_,
		img::bitmap & modifiable_, float max_condition_number_)
		: src(src_), matte(matte_), next_matte(next_matte_), modifiable(modifiable_),
		  max_condition_number(max_condition_number_)
	{
		clear(); 
		mv = gsl_matrix_view_array(&m[0][0], 6, 6);
		Vv = gsl_matrix_view_array(&V[0][0], 6, 6);
		Sv = gsl_vector_view_array(S, 6);
		work_v = gsl_vector_view_array(work, 6);
		bv = gsl_vector_view_array(b, 6);
		gv = gsl_vector_view_array(g, 6);
		rv = gsl_vector_view_array(r, 6);
	}
	void clear() {
		remove_all();
		checks = achanges = 0;
	}
	size_extent dim() const
		{ return src.dim(); }
	uint check_count() const
		{ return checks; }
	uint change_count() const
		{ return achanges; }
	bool selected(size_point const & p) const
		{ return modifiable[p]; }
	void add(size_point const & p)
		{ add_sample(v(src[p]), p.x, p.y, matte[p]); }
	void remove(size_point const & p)
		{ remove_sample(v(src[p]), p.x, p.y, matte[p]); }
	void remove_all() {
		n = A = AA = X = Y = XX = XY = YY = XA = YA = 0;
		XAA = YAA = XXA = XYA = YYA = 0;
		XXAA = XYAA = YYAA = 0;
		I = AI = XI = YI = vectoru::nill;
		AXI = AYI = vectorul::nill;
	}
	void process(size_point const & p, size_rect const &) {
		++checks;
		grey_mask_value g;
		if (!classify(v(src[p]), p, g) || matte[p] == g)
			modifiable[p] = false;
		else {
			++achanges;
			next_matte[p] = g;
		}
	}

private:
	/// guesses alpha for the given color
	bool classify(vectorf const & c, size_point const & p, grey_mask_value & gmv) const 
	{
		//         [ a b ]
		// A^T A = [ b c ]
		//float aa = AA;
		//float bb = uchar_max * A - AA;
		//float cc = uchar_max * uchar_max * n + AA - 2 * uchar_max * A;
		//float discriminant = sqr(aa - cc) + 4 * bb * bb;
		//float condition_number = ((aa + cc) + sqrt(discriminant)) / ((aa + cc) - sqrt(discriminant));
		//if (sqr(max_condition_number + 1) * discriminant > sqr((max_condition_number - 1) * (a + c)))
		//	return false;

		m[0][0] = AA;
		m[0][1] = m[1][0] = XAA;
		m[0][2] = m[2][0] = YAA;
		m[0][3] = m[3][0] = uchar_max * A - AA;
		m[0][4] = m[4][0] = uchar_max * double(XA) - XAA;
		m[0][5] = m[5][0] = uchar_max * double(YA) - YAA;
		m[1][1] = XXAA;
		m[1][2] = m[2][1] = XYAA;
		m[1][3] = m[3][1] = uchar_max * double(XA) - XAA;
		m[1][4] = m[4][1] = uchar_max * double(XXA) - XXAA;
		m[1][5] = m[5][1] = uchar_max * double(XYA) - XYAA;
		m[2][2] = YYAA;
		m[2][3] = m[3][2] = uchar_max * double(YA) - YAA;
		m[2][4] = m[4][2] = uchar_max * double(XYA) - XYAA;
		m[2][5] = m[5][2] = uchar_max * double(YYA) - YYAA;
		m[3][3] = uchar_max * uchar_max * n + AA - 2 * uchar_max * A;
		m[3][4] = m[4][3] = double(uchar_max * uchar_max) * X - 2 * uchar_max * double(XA) + XAA;
		m[3][5] = m[5][3] = double(uchar_max * uchar_max) * Y - 2 * uchar_max * double(YA) + YAA;
		m[4][4] = double(uchar_max * uchar_max) * XX - 2 * uchar_max * double(XXA) + XXAA;
		m[4][5] = m[5][4] = double(uchar_max * uchar_max) * XY - 2 * uchar_max * double(XYA) + XYAA;
		m[5][5] = double(uchar_max * uchar_max) * YY - 2 * uchar_max * double(YYA) + YYAA;

		//compute: m = U * S * V^T; U replaces m
		gsl_linalg_SV_decomp(&mv.matrix, &Vv.matrix, &Sv.vector, &work_v.vector);
		if (S[5] * max_condition_number < S[4])
			return false;

		work[0] = double(uchar_max) * AI.x;
		work[1] = double(uchar_max) * AXI.x;
		work[2] = double(uchar_max) * AYI.x;
		work[3] = double(uchar_max) * (uchar_max * I.x - AI.x);
		work[4] = double(uchar_max) * (uchar_max * XI.x - AXI.x);
		work[5] = double(uchar_max) * (uchar_max * YI.x - AYI.x);
		gsl_linalg_SV_solve(&mv.matrix, &Vv.matrix, &Sv.vector, &work_v.vector, &bv.vector);
		
		work[0] = double(uchar_max) * AI.y;
		work[1] = double(uchar_max) * AXI.y;
		work[2] = double(uchar_max) * AYI.y;
		work[3] = double(uchar_max) * (uchar_max * I.y - AI.y);
		work[4] = double(uchar_max) * (uchar_max * XI.y - AXI.y);
		work[5] = double(uchar_max) * (uchar_max * YI.y - AYI.y);
		gsl_linalg_SV_solve(&mv.matrix, &Vv.matrix, &Sv.vector, &work_v.vector, &gv.vector);

		work[0] = double(uchar_max) * AI.z;
		work[1] = double(uchar_max) * AXI.z;
		work[2] = double(uchar_max) * AYI.z;
		work[3] = double(uchar_max) * (uchar_max * I.z - AI.z);
		work[4] = double(uchar_max) * (uchar_max * XI.z - AXI.z);
		work[5] = double(uchar_max) * (uchar_max * YI.z - AYI.z);
		gsl_linalg_SV_solve(&mv.matrix, &Vv.matrix, &Sv.vector, &work_v.vector, &rv.vector);

		vectorf F0(b[0], g[0], r[0]);
		vectorf Fx(b[1], g[1], r[1]);
		vectorf Fy(b[2], g[2], r[2]);
		vectorf B0(b[3], g[3], r[3]);
		vectorf Bx(b[4], g[4], r[4]);
		vectorf By(b[5], g[5], r[5]);

		/* extrapolation never works :-(
		vectorf Fn = outer_product(Fx, Fy);
		float Fn_sqr = sqr(Fn);
		vectorf Bn = outer_product(Bx, By);
		float Bn_sqr = sqr(Bn);
		if (Fn_sqr > std::numeric_limits<float>::epsilon() &&
			Bn_sqr > std::numeric_limits<float>::epsilon())
		{
			//both foreground and background colors are distributed over planes
			Fn *= 1 / sqrt(Fn_sqr);
			Bn *= 1 / sqrt(Bn_sqr);
			if (scalar_product(Fn, Bn) < 0)
				Bn = -Bn;

			//signed distance of the current color to the foreground plane
			//(negative if point is from one side and positive if from another)
			float dF = scalar_product(c - F0, Fn);
			float dB = scalar_product(c - B0, Bn);
			float hF = fabs(dF);
			float hB = fabs(dB);
			if (dF * dB < 0)
				gmv = uchar_max * (hB / (hF + hB)) + 0.5f;
			else if (hF < hB)
				gmv = uchar_max;
			else
				gmv = 0;
		}
		else {*/
			//either foreground or background colors are not distributed over a plane;
			//assume that both of them are monochrome
			vectorf F = F0 + float(p.x) * Fx + float(p.y) * Fy;
			vectorf B = B0 + float(p.x) * Bx + float(p.y) * By;
			vectorf D = F - B;

			float den = sqr(D);
			if (den == 0)
				gmv = uchar_avg;
			else {
				float num = scalar_product(c - B, D);
				if (num <= 0)
					gmv = 0;
				else if (num >= den)
					gmv = uchar_max;
				else 
					gmv = uchar_max * num / den + 0.5f;
			}
		//}
		return true;
	}
private:
	void add_sample(vectoru const & c, uint x, uint y, uint a) {
		++n;
		A += a;
		AA += a * a;
		X += x;
		Y += y;
		XX += x * x;
		XY += x * y;
		YY += y * y;
		XA += x * a;
		YA += y * a;
		XAA += x * a * a;
		YAA += y * a * a;
		XXA += x * x * a;
		XYA += x * y * a;
		YYA += y * y * a;
		XXAA += x * ulong(x * a * a);
		XYAA += x * ulong(y * a * a);
		YYAA += y * ulong(y * a * a);
		I += c;
		AI += a * c;
		XI += x * c;
		YI += y * c;
		AXI += a * x * c;
		AYI += a * y * c;
	}
	void remove_sample(vectoru const & c, uint x, uint y, uint a) {
		--n;
		A -= a;
		AA -= a * a;
		X -= x;
		Y -= y;
		XX -= x * x;
		XY -= x * y;
		YY -= y * y;
		XA -= x * a;
		YA -= y * a;
		XAA -= x * a * a;
		YAA -= y * a * a;
		XXA -= x * x * a;
		XYA -= x * y * a;
		YYA -= y * y * a;
		XXAA -= x * ulong(x * a * a);
		XYAA -= x * ulong(y * a * a);
		YYAA -= y * ulong(y * a * a);
		I -= c;
		AI -= a * c;
		XI -= x * c;
		YI -= y * c;
		AXI -= a * x * c;
		AYI -= a * y * c;
	}
private:
	Raster const & src;
    Matte const & matte;
	grey_mask::mask & next_matte;
	img::bitmap & modifiable;
	float max_condition_number;
private:
	uint n, A, AA;
	uint X, Y, XX, XY, YY, XA, YA;
	ulong XAA, YAA;
	ulong XXA, XYA, YYA;
	ulong XXAA, XYAA, YYAA;
	vectoru I, AI, XI, YI;
	vectorul AXI, AYI;
	uint checks, achanges;
private:
	mutable double m[6][6], V[6][6], S[6], work[6], b[6], g[6], r[6];
	mutable gsl_matrix_view mv, Vv;
	mutable gsl_vector_view Sv, work_v, bv, gv, rv;
};

template <typename Raster, typename Matte>
static void perform_core(
	const Raster & src, 
	Matte & matte, 
	const img::bitmap & strict_fore, 
	const img::bitmap & strict_back,
	img::bitmap & modifiable, ///< map of pixels that may be changed during the following pass
	const params & p)
{
	if (!p.silent)
		clog << "Volosok(R) matting..." << endl;

	size_extent dim = src.dim();
	assert (dim == matte.dim());
	assert (dim == modifiable.dim());
	int_rect frame(0, 0, dim.sx, dim.sy);

	grey_mask::mask next_matte = matte;

	for (uint pass = 0; pass < p.passes_num; ++pass) {
		if (!p.silent)
			clog << "pass " << pass+1 << '/' << p.passes_num << flush;

		img::combine(modifiable, strict_fore, img::bit_and_not());
		img::combine(modifiable, strict_back, img::bit_and_not());
		smoother<Raster, Matte> sm(src, matte, next_matte, modifiable, p.max_condition_number);
		window_pass(sm, p.whs);

		if (!p.silent)
			clog << ": checked = " << sm.check_count() << ", alpha changes = " << sm.change_count() << endl;

		if (sm.change_count() == 0)
			break;

		matte = next_matte;

		img::dilate8(modifiable, p.whs);
	}
}

template <typename Raster, typename Matte>
static void perform_from_bit_mask(
	const Raster & src, 
	const bit_mask::mask & m, 
	Matte & matte, 
	const params & p)
{
	size_extent dim = src.dim();
	assert (dim == m.dim());
	assert (dim == matte.dim());

	for (typename Matte::iterator i = matte.begin(); i != matte.end(); ++i)
		*i = m[i] ? uchar_max : 0;

	img::bitmap strict_fore, strict_back;
	img::erode4(strict_fore, m, p.fore_depth);
	img::dilate4(strict_back, m, p.back_depth);
	strict_back.inverse();

	img::bitmap modifiable = m; 
	{
		uint w = 1 + p.whs * (p.max_condition_number - 1) / (p.max_condition_number + 1);
		img::dilate4(modifiable, w);
		img::bitmap tmp = m;
		tmp.inverse();
		img::dilate4(tmp, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(src, matte, strict_fore, strict_back, modifiable, p);
}

template <typename Raster, typename Matte>
static void perform_from_grey_mask(
	const Raster & src, 
	Matte & matte, 
	const params & p)
{
	size_extent dim = matte.dim();

	img::bitmap not_min(dim), not_max(dim), strict_fore(dim), strict_back(dim);
	for (typename Matte::/*const_*/iterator i = matte.begin(); i != matte.end(); ++i) {
		not_min[i] = *i != 0;
		not_max[i] = *i != uchar_max;
		strict_fore[i] = *i > uchar_avg;
	}

	strict_back = strict_fore;
	img::erode4(strict_fore, p.fore_depth);
	img::dilate4(strict_back, p.back_depth);
	strict_back.inverse();

	img::bitmap modifiable;
	{
		uint w = 1 + p.whs * (p.max_condition_number - 1) / (p.max_condition_number + 1);
		img::dilate4(modifiable, not_min, w);
		img::bitmap tmp;
		img::dilate4(tmp, not_max, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(src, matte, strict_fore, strict_back, modifiable, p);
}

// exported functions

void perform(
	const byte_image & src, 
	const bit_mask::mask & m, 
	grey_mask::mask & matte, 
	const params & p)
{
	matte.resize(src.dim());
	perform_from_bit_mask(src, m, matte, p);
}

void perform(
	const byte_image & src, 
	grey_mask::mask & matte, 
	const params & p)
{
	matte.resize(src.dim());
	perform_from_grey_mask(src, matte, p);
}

void perform(
	byte_image & im, 
	const params & p)
{
	perform_from_grey_mask(im, img::projection(im, &byte_color::a), p);
}

void perform(
	argb_byte_image & im, 
	const params & p)
{
	perform_from_grey_mask(im, img::projection(im, &argb_byte_color::a), p);
}

} //namespace smooth_d

} //namespace refine
